/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: 集合通信算子kernel
 * Author: lilianlin
 * Create: 2022-07-15
 */

#include "hcom_launch_kernel.h"
#include <hccl.h>
#include "hccl/hcom.h"
#include "hcom_pub.h"
#include "workflow_pub.h"
#include "../op_base/src/op_base.h"
#include "hcom_ops_kernel_info_store.h"
#include "exe_graph/runtime/gert_mem_allocator.h"
#include "exe_graph/runtime/gert_tensor_data.h"
#include "hcom_ops_kernel_info_store.h"
#include "comm.h"
#include "adapter/adapter_rts.h"
#include "op_base_rpc.h"

namespace hccl {
enum class HcomOpLaunchInput {
    OP_ARGS,
    STREAM,
    INPUT_NUM,
    OUTPUT_NUM,
};

enum class RecvOpLaunchInput {
    OP_ARGS,
    STREAM,
    ALLOCATOR,
    DATATYPE,
};

enum RemoteLookupInputIndex {
    LOOKUP_IT_KEY_IDX = 0,
    LOOKUP_IT_TABLEID_IDX = 1,
    LOOKUP_IT_IDX_NUM = 2,
};

enum RemoteLookupOutputIndex {
    LOOKUP_OT_VALUE_IDX = 0,
    LOOKUP_OT_IDX_NUM = 1,
};

uint64_t RoundUp(const uint64_t originValue, const uint64_t multiple)
{
    if (multiple == 0) {
        return 0;
    }
    return (originValue + multiple - 1) / multiple * multiple;
}

HcclResult GetCountByShapeAlign(const gert::Shape& shape, HcclDataType dataType, uint64_t& count)
{
    const uint64_t memAlignSize = 512;
    uint32_t unitSize = SIZE_TABLE[dataType];
    count = RoundUp(shape.GetShapeSize() * unitSize, memAlignSize) / unitSize;
    HCCL_INFO("[debug][print]shape size:%ld, count:%lu", shape.GetShapeSize(), count);
    return HCCL_SUCCESS;
}

HcclResult GetSendCountByShapeAlign(const gert::Shape& shape, HcclDataType dataType, uint64_t& count)
{
    const uint64_t memAlignSize = 32;
    uint32_t unitSize = SIZE_TABLE[dataType];
    count = RoundUp(shape.GetShapeSize() * unitSize, memAlignSize) / unitSize;
    HCCL_INFO("[debug][print]shape size:%ld, count:%lu", shape.GetShapeSize(), count);
    return HCCL_SUCCESS;
}

HcclResult GetCountByShape(const gert::Shape& shape, HcclDataType dataType, uint64_t& count)
{
    count = shape.GetShapeSize();
    return HCCL_SUCCESS;
}

HcclResult RegisterAddInfosFromLaunchArgs(HcomOpLaunchArgs& launchArgs, vector<HcomRemoteAccessAddrInfo>& addrInfos)
{
    // 把infer shape中的tensor放入vector中
    gert::Tensor* tensor = static_cast<gert::Tensor*>(launchArgs.inputAddrs[0]);
    uint64_t* data = tensor->GetData<uint64_t>();
    gert::Tensor* local_tensor = static_cast<gert::Tensor*>(launchArgs.inputAddrs[1]);
    uint64_t* local_data = local_tensor->GetData<uint64_t>();
    const int64_t dims = launchArgs.inputShapes[0].GetDim(0);
    addrInfos.resize(dims);
    for (int64_t idx = 0UL; idx < dims; idx++) {
        int64_t idx_line = idx * INDEX_PARA_3;
        addrInfos[idx] = {static_cast<uint32_t>(data[idx_line]),
                          data[idx_line + 1],
                          ge::PtrToValue(local_data),
                          data[idx_line + INDEX_PARA_2]};
    }
    return HCCL_SUCCESS;
}

HcclResult ExtractRemoteReadAddrInfos(HcomOpLaunchArgs& launchArgs, vector<HcomRemoteAccessAddrInfo>& addrInfos)
{
    u64* remoteData = static_cast<u64*>(launchArgs.inputAddrs[0]);
    u64 localData = reinterpret_cast<u64>(launchArgs.outputAddrs[0]);

    const int64_t rowNum = launchArgs.inputShapes[0].GetDim(0);
    addrInfos.reserve(rowNum);
    for (int64_t idx = 0UL; idx < rowNum; idx++) {
        int64_t lineIdx = idx * INDEX_PARA_3;
        if (remoteData[lineIdx + 1] == 0xFFFFFFFFFFFFFFFF) {
            break;
        }
        addrInfos.push_back({static_cast<u32>(remoteData[lineIdx]),
                          remoteData[lineIdx + 1],
                          localData,
                          remoteData[lineIdx + INDEX_PARA_2]});
        localData += remoteData[lineIdx + INDEX_PARA_2];
    }
    return HCCL_SUCCESS;
}

HcclResult ExtractRemoteAddrInfosWithOffset(HcomOpLaunchArgs& launchArgs, vector<HcomRemoteAccessAddrInfo>& addrInfos)
{
    u64* remoteData = static_cast<u64*>(launchArgs.inputAddrs[0]);
    u64 localData = reinterpret_cast<u64>(launchArgs.inputAddrs[1]);
    u64* offsetData = static_cast<u64*>(launchArgs.inputAddrs[2]);  // 输入index 2为offset

    const int64_t rowNum = launchArgs.inputShapes[0].GetDim(0);
    addrInfos.reserve(rowNum);
    for (int64_t idx = 0UL; idx < rowNum; idx++) {
        int64_t lineIdx = idx * INDEX_PARA_3;
        if (remoteData[lineIdx + 1] == 0xFFFFFFFFFFFFFFFF) {
            break;
        }
        addrInfos.push_back({static_cast<u32>(remoteData[lineIdx]),
                          remoteData[lineIdx + 1],
                          localData + offsetData[idx],
                          remoteData[lineIdx + 2]});
    }
    return HCCL_SUCCESS;
}

HcclResult ExtractRemoteWriteAddsInfos(HcomOpLaunchArgs& launchArgs, vector<HcomRemoteAccessAddrInfo>& addrInfos)
{
    u64* remoteData = static_cast<u64*>(launchArgs.inputAddrs[0]);
    u64 localData = reinterpret_cast<u64>(launchArgs.inputAddrs[1]);

    const int64_t rowNum = launchArgs.inputShapes[0].GetDim(0);
    addrInfos.reserve(rowNum);
    for (int64_t idx = 0UL; idx < rowNum; idx++) {
        int64_t lineIdx = idx * INDEX_PARA_3;
        if (remoteData[lineIdx + 1] == 0xFFFFFFFFFFFFFFFF) {
            break;
        }
        addrInfos.push_back({static_cast<u32>(remoteData[lineIdx]),
                          remoteData[lineIdx + 1],
                          localData,
                          remoteData[lineIdx + INDEX_PARA_2]});
        localData += remoteData[lineIdx + INDEX_PARA_2];
    }
    return HCCL_SUCCESS;
}

HcclResult ExtractRemoteScatterWriteAddrInfos(HcomOpLaunchArgs& launchArgs, vector<HcomRemoteAccessAddrInfo>& addrInfos)
{
    if (launchArgs.inputNum == 3) {     // input个数为3时，输入包含offset
        return ExtractRemoteAddrInfosWithOffset(launchArgs, addrInfos);
    } else if (launchArgs.inputNum == 2) {  // input个数为2时，输入不包含offset
        return ExtractRemoteWriteAddsInfos(launchArgs, addrInfos);
    }
    HCCL_ERROR("HcomRemoteScatterWrite: input num[%u] is invalid.", launchArgs.inputNum);
    return HCCL_E_PARA;
}

HcclResult HcomAllGatherKernel(HcomOpLaunchArgs& launchArgs)
{
    std::shared_ptr<hcclComm> hcclComm;
    CHK_RET(HcomGetCommByGroup(launchArgs.opAttr.group, hcclComm));
    CHK_PRT_RET((launchArgs.inputNum != 1), HCCL_ERROR("HcomAllGather input num[%u] is invalid.", launchArgs.inputNum),
        HCCL_E_PARA);

    uint64_t count;
    CHK_RET(GetCountByShape(launchArgs.inputShapes[0], launchArgs.opAttr.dataType, count));
    CHK_RET(HcclAllGather(launchArgs.inputAddrs[0], launchArgs.outputAddrs[0], count,
        launchArgs.opAttr.dataType, static_cast<HcclComm>(hcclComm.get()), launchArgs.stream));

    return HCCL_SUCCESS;
}

HcclResult HcomAllGatherVKernel(HcomOpLaunchArgs& launchArgs)
{
    std::shared_ptr<hcclComm> hcclComm;
    CHK_RET(HcomGetCommByGroup(launchArgs.opAttr.group, hcclComm));
    uint64_t sendCount;
    CHK_RET(GetCountByShape(launchArgs.inputShapes[0], launchArgs.opAttr.dataType, sendCount));
    u32 rankSize = 0;
    hcclComm->GetRankSize(rankSize);
    uint64_t dimzero = static_cast<uint64_t>(launchArgs.inputShapes[0].GetDim(0));
    uint64_t multiplier = (dimzero == 0) ? 1 : (sendCount / dimzero);
    uint64_t* recvcounts = static_cast<uint64_t *>(const_cast<void*>(launchArgs.inputAddrs[2]));
    std::vector<uint64_t> recvDispls;
    uint64_t tempSum = 0;
    for (u32 i = 0; i < rankSize; i++) {
        recvcounts[i] *= multiplier;
        recvDispls.push_back(tempSum);
        tempSum += recvcounts[i];
    }
    const void* recvDisplsPtr = static_cast<const void*>(recvDispls.data());
    const void* recvcountsPtr = static_cast<const void*>(recvcounts);
    CHK_RET(HcclAllGatherV(launchArgs.inputAddrs[0], sendCount, launchArgs.outputAddrs[0],
        recvcountsPtr, recvDisplsPtr, launchArgs.opAttr.dataType, static_cast<HcclComm>(hcclComm.get()),
        launchArgs.stream));
    return HCCL_SUCCESS;
}

HcclResult HcomGetBatchAllReduceInfo(const HcomOpLaunchArgs &launchArgs, uint64_t maxCommCount,
    std::vector<uint32_t> &sliceIdxs, std::vector<uint64_t> &inputsCount)
{
    uint32_t sliceIdx = 0;
    uint32_t sliceSum = 0;
    uint64_t commCount = 0;
    inputsCount.resize(launchArgs.inputNum);
    for (uint32_t i = 0; i < launchArgs.inputNum; i++) {
        uint64_t count;
        CHK_RET(GetCountByShape(launchArgs.inputShapes[i], launchArgs.opAttr.dataType, count));
        HCCL_DEBUG("HcomAllReduceKernel: input[%u] count: %llu", i, count);
        uint64_t outCount;
        CHK_RET(GetCountByShape(launchArgs.outputShapes[i], launchArgs.opAttr.dataType, outCount));
        CHK_PRT_RET((count != outCount),
            HCCL_ERROR("HcomAllReduceKernel: input[%u] count: %llu is different with output[%u] count: %llu", i, count,
            i, outCount),
            HCCL_E_PARA);

        commCount += count;
        inputsCount[i] = count;
        if (commCount <= maxCommCount) {
            sliceIdx++;
        } else {
            if (sliceIdx == 0) {
                sliceIdx++;
                HCCL_DEBUG("HcomAllReduceKernel: slice[%zu] has %u tensors.", sliceIdxs.size(), sliceIdx);
                sliceIdxs.push_back(sliceIdx);
                sliceSum += sliceIdx;
                commCount = 0;
                sliceIdx = 0;
            } else {
                HCCL_DEBUG("HcomAllReduceKernel: slice[%zu] has %u tensors.", sliceIdxs.size(), sliceIdx);
                sliceIdxs.push_back(sliceIdx);
                sliceSum += sliceIdx;
                commCount = count;
                sliceIdx = 1;
            }
        }
    }

    if (sliceSum < launchArgs.inputNum) {
        HCCL_DEBUG("HcomAllReduceKernel: slice[%zu] has %u tensors.", sliceIdxs.size(),
            (launchArgs.inputNum - sliceSum));
        sliceIdxs.push_back(launchArgs.inputNum - sliceSum);
    }
    return HCCL_SUCCESS;
}

// 将从inputsOffset开始的连续inputsNum个通信算子输入tensor copy至CCL buffer，并输出copy完成的数据量
HcclResult HcomCopyInputsToCCLbuff(const HcomOpLaunchArgs &launchArgs, uint32_t inputsNum, uint32_t inputsOffset,
    std::vector<uint64_t> &inputsCount, void *cclBuff, uint64_t cclBuffSize, uint64_t &commCount)
{
    uint32_t unitSize = SIZE_TABLE[launchArgs.opAttr.dataType];
    u8 *dstBuff = static_cast<u8 *>(cclBuff);
    uint64_t maxDstSize = cclBuffSize;
    for (uint32_t j = 0; j < inputsNum; j++) {
        uint32_t inputIdx = inputsOffset + j;
        CHK_RET(hrtMemAsyncCopyByQos(dstBuff, maxDstSize, launchArgs.inputAddrs[inputIdx],
            inputsCount[inputIdx] * unitSize, HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_DEVICE_TO_DEVICE, launchArgs.stream,
            launchArgs.opAttr.qosConfig));
        HCCL_DEBUG("HcomAllReduceKernel: copy input[%u] addr[%p] len[%u] to cclbuffer[%p].", inputIdx,
            launchArgs.inputAddrs[inputIdx], inputsCount[inputIdx] * unitSize, dstBuff);
        dstBuff += inputsCount[inputIdx] * unitSize;
        maxDstSize -= inputsCount[inputIdx] * unitSize;
        commCount += inputsCount[inputIdx];
    }
    return HCCL_SUCCESS;
}

// 将CCL buffer中的数据copy至从outputsOffset开始的连续outputsNum个通信算子输出tensor
HcclResult HcomCopyCCLbuffToOutnputs(const HcomOpLaunchArgs &launchArgs, uint32_t outputsNum, uint32_t outputsOffset,
    std::vector<uint64_t> &outputsCount, void *cclBuff)
{
    uint32_t unitSize = SIZE_TABLE[launchArgs.opAttr.dataType];
    u8 *srcBuff = static_cast<u8 *>(cclBuff);
    for (uint32_t j = 0; j < outputsNum; j++) {
        uint32_t outputIdx = outputsOffset + j;
        CHK_RET(hrtMemAsyncCopyByQos(launchArgs.outputAddrs[outputIdx], outputsCount[outputIdx] * unitSize, srcBuff,
            outputsCount[outputIdx] * unitSize, HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_DEVICE_TO_DEVICE,
            launchArgs.stream, launchArgs.opAttr.qosConfig));
        HCCL_DEBUG("HcomAllReduceKernel: copy cclbuffer[%p] to output[%u] addr[%p] len[%u].", srcBuff, outputIdx,
            launchArgs.outputAddrs[outputIdx], outputsCount[outputIdx] * unitSize);
        srcBuff += outputsCount[outputIdx] * unitSize;
    }
    return HCCL_SUCCESS;
}

HcclResult HcomAllReduceKernel(HcomOpLaunchArgs& launchArgs)
{
    std::shared_ptr<hcclComm> hcclComm;
    CHK_RET(HcomGetCommByGroup(launchArgs.opAttr.group, hcclComm));
    CHK_PRT_RET((launchArgs.inputNum < 1), HCCL_ERROR("HcomAllReduce input num[%u] is invalid.", launchArgs.inputNum),
        HCCL_E_PARA);

    void *commInputPtr = nullptr;
    u64 commInputSize = 0;
    void *commOutputPtr = nullptr;
    u64 commOutputSize = 0;
    CHK_RET(hcclComm->GetInCCLbuffer(commInputPtr, commInputSize));
    CHK_RET(hcclComm->GetOutCCLbuffer(commOutputPtr, commOutputSize));
    uint64_t maxCommCount = commInputSize / SIZE_TABLE[launchArgs.opAttr.dataType];

    std::vector<uint32_t> sliceIdxs; // 按照cclbuffer size将输入tensors分段，每一段的tensor个数
    std::vector<uint64_t> inputsCount;
    CHK_RET(HcomGetBatchAllReduceInfo(launchArgs, maxCommCount, sliceIdxs, inputsCount));
    uint32_t inputOffset = 0;
    for (uint32_t i = 0; i < sliceIdxs.size(); i++) {
        if (sliceIdxs[i] == 1) {
            u8 *curInputPtr = static_cast<u8 *>(launchArgs.inputAddrs[inputOffset]);
            u8 *curOutputPtr = static_cast<u8 *>(launchArgs.outputAddrs[inputOffset]);

            CHK_RET(HcclAllReduce(curInputPtr, curOutputPtr, inputsCount[inputOffset], launchArgs.opAttr.dataType,
                    launchArgs.opAttr.op.allreduce.reduction, static_cast<HcclComm>(hcclComm.get()),
                    launchArgs.stream));
        } else {
            uint64_t commCount = 0;
            // step1 将allreduce的输入tensor copy到in cclbuffer
            CHK_RET(HcomCopyInputsToCCLbuff(launchArgs, sliceIdxs[i], inputOffset, inputsCount, commInputPtr,
                commInputSize, commCount));

            // step2 执行allreduce
            CHK_RET(HcclAllReduce(commInputPtr, commOutputPtr, commCount, launchArgs.opAttr.dataType,
                launchArgs.opAttr.op.allreduce.reduction, static_cast<HcclComm>(hcclComm.get()), launchArgs.stream));

            // step3 将out cclbuffer的内容copy到allreduce的输出tensor
            CHK_RET(HcomCopyCCLbuffToOutnputs(launchArgs, sliceIdxs[i], inputOffset, inputsCount, commOutputPtr));
        }

        inputOffset += sliceIdxs[i];
    }
    return HCCL_SUCCESS;
}

HcclResult HcomBroadcastKernel(HcomOpLaunchArgs& launchArgs)
{
    std::shared_ptr<hcclComm> hcclComm;
    CHK_RET(HcomGetCommByGroup(launchArgs.opAttr.group, hcclComm));
    CHK_PRT_RET((launchArgs.inputNum < 1), HCCL_ERROR("HcomBroadcast input num[%u] is invalid.", launchArgs.inputNum),
        HCCL_E_PARA);

    uint64_t count;
    CHK_RET(GetCountByShapeAlign(launchArgs.inputShapes[0], launchArgs.opAttr.dataType, count));
    CHK_RET(HcclBroadcast(launchArgs.inputAddrs[0], count, launchArgs.opAttr.dataType,
        launchArgs.opAttr.op.broadcast.root, static_cast<HcclComm>(hcclComm.get()), launchArgs.stream));
    return HCCL_SUCCESS;
}

HcclResult HcomReduceScatterKernel(HcomOpLaunchArgs& launchArgs)
{
    std::shared_ptr<hcclComm> hcclComm;
    CHK_RET(HcomGetCommByGroup(launchArgs.opAttr.group, hcclComm));
    CHK_PRT_RET((launchArgs.inputNum != 1), HCCL_ERROR("HcomReduceScatter input num[%u] is invalid.",
        launchArgs.inputNum), HCCL_E_PARA);

    uint64_t count;
    CHK_RET(GetCountByShape(launchArgs.outputShapes[0], launchArgs.opAttr.dataType, count));
    CHK_RET(HcclReduceScatter(launchArgs.inputAddrs[0], launchArgs.outputAddrs[0], count, launchArgs.opAttr.dataType,
        launchArgs.opAttr.op.reducescatter.reduction, static_cast<HcclComm>(hcclComm.get()), launchArgs.stream));
    return HCCL_SUCCESS;
}

HcclResult HcomReduceScatterVKernel(HcomOpLaunchArgs& launchArgs)
{
    std::shared_ptr<hcclComm> hcclComm;
    CHK_RET(HcomGetCommByGroup(launchArgs.opAttr.group, hcclComm));

    uint64_t recvCount;
    CHK_RET(GetCountByShape(launchArgs.outputShapes[0], launchArgs.opAttr.dataType, recvCount));
    u32 rankSize = 0;
    hcclComm->GetRankSize(rankSize);
    // 计算第0维后的维度乘积
    uint64_t dimzero = static_cast<uint64_t>(launchArgs.outputShapes[0].GetDim(0));
    uint64_t multiplier = (dimzero == 0) ? 1 : (recvCount / dimzero);
    uint64_t* sendCounts = static_cast<uint64_t *>(const_cast<void*>(launchArgs.inputAddrs[INPUT_INDEX_2]));
    vector<int64_t> sendDispls;
    int64_t tmpCount = 0;
    for(size_t i=0; i<rankSize; i++) {
        sendCounts[i] *= multiplier;
        sendDispls.push_back(tmpCount);
        tmpCount += sendCounts[i];
    }
    const void* sendCountsPtr = static_cast<const void*>(sendCounts);
    const void* sendDisplsPtr = static_cast<const void*>(sendDispls.data());

    CHK_RET(HcclReduceScatterV(launchArgs.inputAddrs[0], sendCountsPtr, sendDisplsPtr,
        launchArgs.outputAddrs[0], recvCount,
        launchArgs.opAttr.dataType, launchArgs.opAttr.op.reducescatterv.reduction,
        static_cast<HcclComm>(hcclComm.get()), launchArgs.stream));
    return HCCL_SUCCESS;
}

HcclResult HcomAllToAllVKernel(HcomOpLaunchArgs& launchArgs)
{
    std::shared_ptr<hcclComm> hcclComm;
    CHK_RET(HcomGetCommByGroup(launchArgs.opAttr.group, hcclComm));
    CHK_RET(HcclAlltoAllV(launchArgs.inputAddrs[0], launchArgs.inputAddrs[1], launchArgs.inputAddrs[INPUT_INDEX_2],
        launchArgs.opAttr.dataType, launchArgs.outputAddrs[0], launchArgs.inputAddrs[INPUT_INDEX_3],
        launchArgs.inputAddrs[INPUT_INDEX_4], launchArgs.opAttr.op.alltoallv.recvType,
        static_cast<HcclComm>(hcclComm.get()), launchArgs.stream));
    return HCCL_SUCCESS;
}

HcclResult HcomAllToAllVCKernel(HcomOpLaunchArgs& launchArgs)
{
    std::shared_ptr<hcclComm> hcclComm;
    CHK_RET(HcomGetCommByGroup(launchArgs.opAttr.group, hcclComm));
    CHK_RET(HcclAlltoAllVC(launchArgs.inputAddrs[0], launchArgs.inputAddrs[1], launchArgs.opAttr.dataType,
        launchArgs.outputAddrs[0], launchArgs.opAttr.dataType, static_cast<HcclComm>(hcclComm.get()),
        launchArgs.stream));
    return HCCL_SUCCESS;
}

HcclResult HcomGatherAllToAllVKernel(HcomOpLaunchArgs& launchArgs)
{
    std::shared_ptr<hcclComm> hcclComm;
    CHK_RET(HcomGetCommByGroup(launchArgs.opAttr.group, hcclComm));
    HcomGatherAllToAllVParams paramsInfo;
    paramsInfo.addrInfo = launchArgs.inputAddrs[0];
    paramsInfo.addrInfoCountPerRank = launchArgs.inputAddrs[1];
    paramsInfo.recvcounts = launchArgs.inputAddrs[INPUT_INDEX_2];
    paramsInfo.rdispls = launchArgs.inputAddrs[INPUT_INDEX_3];
    paramsInfo.recvbuf = launchArgs.outputAddrs[0];
    paramsInfo.gatheredbuf = launchArgs.outputAddrs[1];
    paramsInfo.addrLength = launchArgs.opAttr.op.gatheralltoallv.addrLength;
    paramsInfo.recvtype = launchArgs.opAttr.dataType;
    CHK_RET(HcclGatherAlltoAllV(paramsInfo, static_cast<HcclComm>(hcclComm.get()), launchArgs.stream));
    return HCCL_SUCCESS;
}

HcclResult HcomRemoteReadKernel(HcomOpLaunchArgs& launchArgs)
{
    std::shared_ptr<RemoteAccess> remoteAccess_;
    vector<HcomRemoteAccessAddrInfo> addrInfos;

    CHK_RET(HcomGetRemoteAccess(remoteAccess_));
    CHK_RET(ExtractRemoteReadAddrInfos(launchArgs, addrInfos));
    CHK_RET(remoteAccess_->RemoteRead(addrInfos, launchArgs.stream));
    return HCCL_SUCCESS;
}

HcclResult HcomRemoteRefReadKernel(HcomOpLaunchArgs& launchArgs)
{
    std::shared_ptr<RemoteAccess> remoteAccess_;
    vector<HcomRemoteAccessAddrInfo> addrInfos;
    CHK_RET(HcomGetRemoteAccess(remoteAccess_));
    CHK_RET(ExtractRemoteAddrInfosWithOffset(launchArgs, addrInfos));
    CHK_RET(remoteAccess_->RemoteRead(addrInfos, launchArgs.stream));
    return HCCL_SUCCESS;
}

HcclResult HcomRemoteWriteKernel(HcomOpLaunchArgs& launchArgs)
{
    std::shared_ptr<RemoteAccess> remoteAccess_;
    vector<HcomRemoteAccessAddrInfo> addrInfos;
    CHK_RET(HcomGetRemoteAccess(remoteAccess_));
    CHK_RET(ExtractRemoteWriteAddsInfos(launchArgs, addrInfos));
    CHK_RET(remoteAccess_->RemoteWrite(addrInfos, launchArgs.stream));
    return HCCL_SUCCESS;
}

HcclResult RemoteScatterWriteKernel(HcomOpLaunchArgs& launchArgs)
{
    std::shared_ptr<RemoteAccess> remoteAccess_;
    vector<HcomRemoteAccessAddrInfo> addrInfos;
    CHK_RET(HcomGetRemoteAccess(remoteAccess_));
    CHK_RET(ExtractRemoteScatterWriteAddrInfos(launchArgs, addrInfos));
    CHK_RET(remoteAccess_->RemoteWrite(addrInfos, launchArgs.stream));
    return HCCL_SUCCESS;
}

HcclResult HcomSendKernel(HcomOpLaunchArgs &launchArgs)
{
    std::shared_ptr<hcclComm> hcclComm;
    CHK_RET(HcomGetCommByGroup(launchArgs.opAttr.group, hcclComm));
    CHK_PRT_RET((launchArgs.inputNum < 1), HCCL_ERROR("HcomSend input num[%u] is invalid.", launchArgs.inputNum),
                HCCL_E_PARA);

    // host内存
    u32 shapeCount = 1 + launchArgs.inputShapes[0].GetDimNum();
    std::vector<int64_t> sendInfo(shapeCount, 0);
    sendInfo[0] = shapeCount - 1;
    for (u32 i = 1; i < sendInfo.size(); i++) {
        sendInfo[i] = launchArgs.inputShapes[0].GetDim(i - 1);
    }

    // dev内存
    auto deleter = [](void *&dst) {
        if (dst != nullptr) {
            CHK_PRT(hrtFree(dst));
            dst = nullptr;
        }
    };

    void *sendShapeMemPtr = nullptr;
    u32 shapeSize = sizeof(int64_t) * shapeCount;
    u32 totalSize = sizeof(int64_t) * SHAPE_INFO_COUNT;
    CHK_RET(hrtMalloc(&sendShapeMemPtr, totalSize));
    std::unique_ptr<void, decltype(deleter)> sendShapeMemPtrUnique(sendShapeMemPtr, deleter);

    // host2dev拷贝
    CHK_RET(hrtMemSyncCopy(sendShapeMemPtr, shapeSize, sendInfo.data(),
                           shapeSize, HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_HOST_TO_DEVICE));

    // send
    CHK_RET(HcclSend(sendShapeMemPtr, SHAPE_INFO_COUNT, HCCL_DATA_TYPE_INT64, launchArgs.opAttr.op.send.destRank,
        static_cast<HcclComm>(hcclComm.get()), launchArgs.stream));

    CHK_RET(hcclStreamSynchronize(launchArgs.stream));

    uint64_t count = 0;
    CHK_RET(GetSendCountByShapeAlign(launchArgs.inputShapes[0], launchArgs.opAttr.dataType, count));
    CHK_RET(HcclSend(launchArgs.inputAddrs[0], count, launchArgs.opAttr.dataType, launchArgs.opAttr.op.send.destRank,
        static_cast<HcclComm>(hcclComm.get()), launchArgs.stream));

    return HCCL_SUCCESS;
}

HcclResult HcomReduceKernel(HcomOpLaunchArgs &launchArgs)
{
    std::shared_ptr<hcclComm> hcclComm;
    CHK_RET(HcomGetCommByGroup(launchArgs.opAttr.group, hcclComm));
    CHK_PRT_RET((launchArgs.inputNum < 1), HCCL_ERROR("HcomReduce input num[%u] is invalid.", launchArgs.inputNum),
        HCCL_E_PARA);

    uint64_t count;
    CHK_RET(GetCountByShape(launchArgs.inputShapes[0], launchArgs.opAttr.dataType, count));
    CHK_RET(HcclReduce(launchArgs.inputAddrs[0], launchArgs.outputAddrs[0], count, launchArgs.opAttr.dataType,
        launchArgs.opAttr.op.reduce.reduction, launchArgs.opAttr.op.reduce.root,
        static_cast<HcclComm>(hcclComm.get()), launchArgs.stream));

    return HCCL_SUCCESS;
}

HcclResult HcomCollRemotePairedParaCheck(const HcomRemoteOperationParams &params)
{
    CHK_PTR_NULL(params.keyAddr);
    CHK_PTR_NULL(params.value);
    CHK_PTR_NULL(params.tableId);
    CHK_PTR_NULL(params.indices);
    CHK_PTR_NULL(params.numUniqued);
    CHK_PTR_NULL(params.psSeg);
    CHK_PTR_NULL(params.psSegNum);

    return HCCL_SUCCESS;
}

HcclResult CollRemoteLookupPairedKernel(HcomOpLaunchArgs& launchArgs)
{
    if (launchArgs.inputNum < LOOKUP_PRD_IT_IDX_NUM) {
        HCCL_ERROR("CollRemoteLookupPaired input num[%u] is error.", launchArgs.inputNum);
        return HCCL_E_PARA;
    }
    if (launchArgs.outputNum < LOOKUP_PRD_OT_IDX_NUM) {
        HCCL_ERROR("CollRemoteLookupPaired output num[%u] is error.", launchArgs.outputNum);
        return HCCL_E_PARA;
    }

    HcomRemoteOperationParams params{};

    params.tableId = static_cast<s32 *>(launchArgs.inputAddrs[LOOKUP_PRD_IT_TABLEID_IDX]);
    params.keyAddr = static_cast<s64 *>(launchArgs.inputAddrs[LOOKUP_PRD_IT_KEY_IDX]);

    params.value = static_cast<void *>(launchArgs.outputAddrs[LOOKUP_PRD_OT_VALUE_IDX]);
    params.indices = static_cast<void *>(launchArgs.outputAddrs[LOOKUP_PRD_OT_INDICES_IDX]);
    params.numUniqued = static_cast<void *>(launchArgs.outputAddrs[LOOKUP_PRD_OT_NUMUNIQUED_IDX]);
    params.psSeg = static_cast<void *>(launchArgs.outputAddrs[LOOKUP_PRD_OT_PSSEG_IDX]);
    params.psSegNum = static_cast<void *>(launchArgs.outputAddrs[LOOKUP_PRD_OT_PSSEGNUM_IDX]);

    CHK_RET(HcomCollRemotePairedParaCheck(params));

    params.count = launchArgs.opAttr.op.collremotelookuppaired.max_num;
    params.valueDim = launchArgs.opAttr.op.collremotelookuppaired.embedding_dim;
    params.tag = launchArgs.opAttr.op.collremotelookuppaired.tag;
    params.stream = launchArgs.stream;
    params.group = launchArgs.opAttr.group;
    params.insertOption = launchArgs.opAttr.op.collremotelookuppaired.insert_option;
    params.flags = launchArgs.opAttr.op.collremotelookuppaired.flags;

    CHK_PRT(HcclCollRemoteLookupOrUpdatePairedOrUniqued(params, EsOpType::LOOKUP));
    HCCL_RUN_INFO("CollRemoteLookupPairedKernel is success");

    return HCCL_SUCCESS;
}

HcclResult CollRemoteUpdatePairedKernel(HcomOpLaunchArgs& launchArgs)
{
    if (launchArgs.inputNum < UPDATE_PRD_IT_IDX_NUM) {
        HCCL_ERROR("CollRemoteUpdatePaired input num[%u] is error.", launchArgs.inputNum);
        return HCCL_E_PARA;
    }

    HcomRemoteOperationParams params{};

    params.tableId = static_cast<s32 *>(launchArgs.inputAddrs[UPDATE_PRD_IT_TABLEID_IDX]);
    params.keyAddr = static_cast<s64 *>(launchArgs.inputAddrs[UPDATE_PRD_IT_KEY_IDX]);
    params.value = static_cast<void *>(launchArgs.inputAddrs[UPDATE_PRD_IT_VALUE_IDX]);
    params.indices = static_cast<void *>(launchArgs.inputAddrs[UPDATE_PRD_IT_INDICES_IDX]);
    params.numUniqued = static_cast<void *>(launchArgs.inputAddrs[UPDATE_PRD_IT_NUMUNIQUED_IDX]);
    params.psSeg = static_cast<void *>(launchArgs.inputAddrs[UPDATE_PRD_IT_PSSEG_IDX]);
    params.psSegNum = static_cast<void *>(launchArgs.inputAddrs[UPDATE_PRD_IT_PSSEGNUM_IDX]);
    if (launchArgs.inputAddrs.size() > UPDATE_PRD_IT_IDX_NUM) {
        params.globalStepAddr = static_cast<void *>(launchArgs.inputAddrs[UPDATE_PRD_IT_IDX_NUM]);
    }

    CHK_RET(HcomCollRemotePairedParaCheck(params));

    params.count = launchArgs.opAttr.op.collremoteupdatepaired.max_num;
    params.valueDim = launchArgs.opAttr.op.collremoteupdatepaired.embedding_dim;
    params.tag = launchArgs.opAttr.op.collremoteupdatepaired.tag;
    params.stream = launchArgs.stream;
    params.group = launchArgs.opAttr.group;

    CHK_PRT(HcclCollRemoteLookupOrUpdatePairedOrUniqued(params, EsOpType::UPDATE));
    HCCL_RUN_INFO("CollRemoteUpdatePairedKernel is success");

    return HCCL_SUCCESS;
}

HcclResult RemoteLookupKernel(HcomOpLaunchArgs& launchArgs)
{
    if (launchArgs.inputNum < LOOKUP_IT_IDX_NUM) {
        HCCL_ERROR("RemoteLookup input num[%u] is error.", launchArgs.inputNum);
        return HCCL_E_PARA;
    }
 
    rtStream_t stream = launchArgs.stream;
    RemoteLookupPubParams lookupPubInfo;
    s32 embeddingDim = launchArgs.opAttr.op.remotelookup.embedding_dim;
    s32 insertFlag = launchArgs.opAttr.op.remotelookup.insert_option;
    s32 esTag = launchArgs.opAttr.op.remotelookup.tag;
    
    lookupPubInfo.keys = static_cast<s64 *>(launchArgs.inputAddrs[LOOKUP_IT_KEY_IDX]);
    lookupPubInfo.keyMaxNum = launchArgs.opAttr.op.remotelookup.max_num;
    lookupPubInfo.valueItemSize = embeddingDim * sizeof(float);
    lookupPubInfo.value = static_cast<void *>(launchArgs.outputAddrs[LOOKUP_OT_VALUE_IDX]);
    lookupPubInfo.tableId = reinterpret_cast<s32*>(launchArgs.inputAddrs[LOOKUP_IT_TABLEID_IDX]);

    CHK_RET(HcclRemoteLookup(esTag, insertFlag, stream, lookupPubInfo));
    HCCL_RUN_INFO("RemoteLookupKernel is success");

    return HCCL_SUCCESS;
}

HcclResult CollRemoteLookupUniquedPairedKernel(HcomOpLaunchArgs& launchArgs)
{
    // 第五项输入为可选项，根据flags在后文判断
    if (launchArgs.inputNum < (LOOKUP_UNIQ_PRD_IT_IDX_NUM - 1)) {
        HCCL_ERROR("CollRemoteLookupUniquedPaired input num[%u] is error.", launchArgs.inputNum);
        return HCCL_E_PARA;
    }
    if (launchArgs.outputNum < LOOKUP_PRD_OT_IDX_NUM) {
        HCCL_ERROR("CollRemoteLookupUniquedPaired output num[%u] is error.", launchArgs.outputNum);
        return HCCL_E_PARA;
    }

    HcomRemoteOperationParams params{};

    params.tableId = static_cast<s32 *>(launchArgs.inputAddrs[LOOKUP_UNIQ_PRD_IT_TABLEID_IDX]);
    params.keyAddr = static_cast<s64 *>(launchArgs.inputAddrs[LOOKUP_UNIQ_PRD_IT_KEY_IDX]);
    params.keyNumInput = static_cast<s64 *>(launchArgs.inputAddrs[LOOKUP_UNIQ_PRD_IT_KEY_NUM_INPUT_IDX]);
    params.uniqueIndices = static_cast<s32 *>(launchArgs.inputAddrs[LOOKUP_UNIQ_PRD_IT_UNIQUE_INDICES_IDX]);

    if (launchArgs.opAttr.op.collremotelookuppairedandunique.flags == ES_FLAGS_ENABLE_COUNTER) {
        if (launchArgs.inputNum < LOOKUP_UNIQ_PRD_IT_IDX_NUM) {
            HCCL_ERROR("CollRemoteLookupUniquedPaired input num[%u] is error while enabling counter.",
                launchArgs.inputNum);
                return HCCL_E_PARA;
        }
        params.keyCount = static_cast<s32 *>(launchArgs.inputAddrs[LOOKUP_UNIQ_PRD_IT_KEY_COUNT_IDX]);
        CHK_PTR_NULL(params.keyCount);
    }

    params.value = static_cast<void *>(launchArgs.outputAddrs[LOOKUP_PRD_OT_VALUE_IDX]);
    params.indices = static_cast<void *>(launchArgs.outputAddrs[LOOKUP_PRD_OT_INDICES_IDX]);
    params.numUniqued = static_cast<void *>(launchArgs.outputAddrs[LOOKUP_PRD_OT_NUMUNIQUED_IDX]);
    params.psSeg = static_cast<void *>(launchArgs.outputAddrs[LOOKUP_PRD_OT_PSSEG_IDX]);
    params.psSegNum = static_cast<void *>(launchArgs.outputAddrs[LOOKUP_PRD_OT_PSSEGNUM_IDX]);

    CHK_RET(HcomCollRemotePairedParaCheck(params));
    CHK_PTR_NULL(params.keyNumInput);
    CHK_PTR_NULL(params.uniqueIndices);

    params.count = launchArgs.opAttr.op.collremotelookuppairedandunique.max_num;
    params.valueDim = launchArgs.opAttr.op.collremotelookuppairedandunique.embedding_dim;
    params.tag = launchArgs.opAttr.op.collremotelookuppairedandunique.tag;
    params.stream = launchArgs.stream;
    params.group = launchArgs.opAttr.group;
    params.insertOption = launchArgs.opAttr.op.collremotelookuppairedandunique.insert_option;
    params.flags = launchArgs.opAttr.op.collremotelookuppairedandunique.flags;

    CHK_PRT(HcclCollRemoteLookupOrUpdatePairedOrUniqued(params, EsOpType::LOOKUP_UNIQUED));
    HCCL_RUN_INFO("CollRemoteLookupUniquedPaired is success");

    return HCCL_SUCCESS;
}

std::vector<std::function<HcclResult(HcomOpLaunchArgs&)>> HcomOpKernelFuncs = {
    HcomAllGatherKernel,
    HcomAllGatherVKernel,
    HcomAllReduceKernel,
    HcomBroadcastKernel,
    HcomReduceScatterKernel,
    HcomReduceScatterVKernel,
    HcomAllToAllVKernel,
    HcomAllToAllVCKernel,
    HcomGatherAllToAllVKernel,
    HcomRemoteReadKernel,
    HcomRemoteRefReadKernel,
    HcomRemoteWriteKernel,
    RemoteScatterWriteKernel,
    HcomSendKernel,
    HcomReduceKernel,
    CollRemoteLookupPairedKernel,
    CollRemoteUpdatePairedKernel,
    CollRemoteLookupUniquedPairedKernel,
    RemoteLookupKernel
};

HcclResult GetHcomOpLaunchArgs(gert::KernelContext *context, HcomOpLaunchArgs& args)
{
    args.opAttr = *(context->GetInputPointer<struct HcomOpAttr>(static_cast<int32_t>(HcomOpLaunchInput::OP_ARGS)));
    HCCL_INFO("GetHcomOpLaunchArgs: opType:%u", args.opAttr.opType);
    HCCL_INFO("GetHcomOpLaunchArgs: dataType:%s", GetDataTypeEnumStr(args.opAttr.dataType).c_str());
    HCCL_INFO("GetHcomOpLaunchArgs: group:%s", args.opAttr.group);

    args.stream = context->GetInputValue<void *>(static_cast<int32_t>(HcomOpLaunchInput::STREAM));
    args.inputNum = context->GetInputValue<uint32_t>(static_cast<int32_t>(HcomOpLaunchInput::INPUT_NUM));
    args.outputNum = context->GetInputValue<uint32_t>(static_cast<int32_t>(HcomOpLaunchInput::OUTPUT_NUM));
    HCCL_INFO("GetHcomOpLaunchArgs: input_num:%u, output_num:%u", args.inputNum, args.outputNum);

    args.inputAddrs.resize(args.inputNum);
    args.outputAddrs.resize(args.outputNum);
    args.inputShapes.resize(args.inputNum);
    args.outputShapes.resize(args.outputNum);

    uint32_t ioNum = args.inputNum + args.outputNum;
    uint32_t i = 0U;
    auto addrStart = static_cast<int32_t>(HcomOpLaunchInput::OUTPUT_NUM) + 1;
    for (; i < args.inputNum; i++) {
        auto tensorData = context->GetInputValue<gert::TensorData*>(addrStart + i);
        args.inputAddrs[i] = tensorData->GetAddr();
        args.inputShapes[i] = *(context->GetInputPointer<gert::Shape>(addrStart + i + ioNum));
        HCCL_INFO("GetHcomOpLaunchArgs: inputAddrs:%p", args.inputAddrs[i]);
    }

    uint32_t j = 0U;
    for (; j < args.outputNum; j++) {
        auto tensorData = context->GetInputValue<gert::TensorData*>(addrStart + i + j);
        args.outputAddrs[j] = tensorData->GetAddr();
        args.outputShapes[j] = *(context->GetInputPointer<gert::Shape>(addrStart + ioNum + args.inputNum + j));
        HCCL_INFO("GetHcomOpLaunchArgs: outputAddrs:%p", args.outputAddrs[j]);
    }

    return HCCL_SUCCESS;
}

bool IsRemoteOpType(HcomOpType opType)
{
    bool isRemoteOpType =
        (opType == HcomOpType::COLL_REMOTE_LOOK_UP_PAIRED || opType == HcomOpType::COLL_REMOTE_UPDATE_PAIRED ||
        opType == HcomOpType::COLL_REMOTE_LOOK_UP_UNIQUED_PAIRED || opType == HcomOpType::REMOTE_LOOK_UP);

    return isRemoteOpType;
}

/*
 * **********************************************************************
 * 注册kernel函数，提供kernel函数的执行原型
 * **********************************************************************
 */
ge::graphStatus LaunchHcomKernel(gert::KernelContext *context)
{
    // 获取算子输入输出资源
    HcomOpLaunchArgs launchArgs;
    if (GetHcomOpLaunchArgs(context, launchArgs) != HCCL_SUCCESS) {
        HCCL_ERROR("LaunchHcomKernel: get hcom op launch args failed.");
        return ge::GRAPH_FAILED;
    }

    if (!IsRemoteOpType(launchArgs.opAttr.opType)) {
        // 获取通信域资源，
        std::shared_ptr<hcclComm> hcclComm;
        if (HcomGetCommByGroup(launchArgs.opAttr.group, hcclComm)) {
            HCCL_ERROR("LaunchHcomKernel: get hcom comm handle failed. group:%s", launchArgs.opAttr.group);
            return ge::GRAPH_FAILED;
        }
        if (hcclComm->CreateCommCCLbuffer()) {
            HCCL_ERROR("LaunchHcomKernel: create ccl buffer failed. group:%s", launchArgs.opAttr.group);
            return ge::GRAPH_FAILED;
        }

        HCCL_DEBUG("Call LaunchHcomKernel, qosCfg value[%u]", launchArgs.opAttr.qosConfig);
        // 下发QosCofig给Runtime
        CHK_RET(hcclComm->SetQosCfg(launchArgs.opAttr.qosConfig));
    }

    // 设置单算子mode
    HcclWorkflowMode lastWorkflowMode = GetWorkflowMode();
    if (SetWorkflowMode(HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) != HCCL_SUCCESS) {
        HCCL_ERROR("LaunchHcomKernel: config opbase mode failed");
        return ge::GRAPH_FAILED;
    }
    SetLaunchKernelMode(true);

    // 调用算子kernel，实现算子下发
    if (HcomOpKernelFuncs[static_cast<int32_t>(launchArgs.opAttr.opType)](launchArgs) != HCCL_SUCCESS) {
        SetLaunchKernelMode(false);
        return ge::GRAPH_FAILED;
    }

    SetLaunchKernelMode(false);
    // workflowmode状态回置
    if (SetWorkflowMode(lastWorkflowMode) != HCCL_SUCCESS) {
        HCCL_ERROR("LaunchHcomKernel: restore hccl mode failed");
        return ge::GRAPH_FAILED;
    }
    return ge::GRAPH_SUCCESS;
}
REGISTER_KERNEL(LaunchHcomKernel).RunFunc(LaunchHcomKernel);

ge::graphStatus LaunchHcomKernelInitComm(gert::KernelContext *context)
{
    // 获取ge option 初始化hccl world group和group list
    CHK_RET(HcomInitialize());
    CHK_RET(InitGroup());
    return ge::GRAPH_SUCCESS;
}

REGISTER_KERNEL(LaunchHcomKernelInitComm).RunFunc(LaunchHcomKernelInitComm);

ge::graphStatus BuildHcclOutputShapeOutputs(const ge::FastNode *node, gert::KernelContext *context)
{
    (void) node;
    auto chainOfShape = context->GetOutput(0);
    if (chainOfShape == nullptr) {
        HCCL_ERROR("BuildHcclOutputShapeOutputs shape failed");
        return ge::GRAPH_FAILED;
    }
    chainOfShape->SetWithDefaultDeleter(new (std::nothrow) gert::StorageShape());
    auto chainOfTensorData = context->GetOutput(1);
    if (chainOfTensorData == nullptr) {
        HCCL_ERROR("BuildHcclOutputShapeOutputs tensordata failed");
        return ge::GRAPH_FAILED;
    }
    chainOfTensorData->SetWithDefaultDeleter(new (std::nothrow) gert::GertTensorData());

    return ge::GRAPH_SUCCESS;
}

HcclResult GetRecvOpLaunchArgs(gert::KernelContext *context, HcomOpLaunchArgs& args)
{
    // args
    args.opAttr = *(context->GetInputPointer<struct HcomOpAttr>(static_cast<int32_t>(RecvOpLaunchInput::OP_ARGS)));
    HCCL_INFO("GetHcomOpLaunchArgs: opType:%u, dataType:%s, group:%s", args.opAttr.opType,
        GetDataTypeEnumStr(args.opAttr.dataType).c_str(), args.opAttr.group);

    // stream
    args.stream = context->GetInputValue<void *>(static_cast<int32_t>(RecvOpLaunchInput::STREAM));

    return HCCL_SUCCESS;
}

ge::graphStatus HcomGetRecvBeforeKernel(HcomOpLaunchArgs& args, std::vector<int64_t>& recvShape)
{
    // 获取通信域资源
    std::shared_ptr<hcclComm> hcclComm;
    if (HcomGetCommByGroup(args.opAttr.group, hcclComm)) {
        HCCL_ERROR("LaunchHcomKernel: get hcom comm handle failed. group:%s", args.opAttr.group);
        return ge::GRAPH_FAILED;
    }

    // 1、create device mem
    auto deleter = [](void *&dst) {
        if (dst != nullptr) {
            CHK_PRT(hrtFree(dst));
            dst = nullptr;
        }
    };

    void *recvShapeMemPtr = nullptr;
    u32 recvShapeSize = sizeof(int64_t) * SHAPE_INFO_COUNT;
    CHK_RET(hrtMalloc(&recvShapeMemPtr, recvShapeSize));
    std::unique_ptr<void, decltype(deleter)> recvShapeMemPtrUnique(recvShapeMemPtr, deleter);

    // 2、recv
    CHK_RET(HcclRecv(recvShapeMemPtr, SHAPE_INFO_COUNT, HCCL_DATA_TYPE_INT64, args.opAttr.op.recv.srcRank,
        static_cast<HcclComm>(hcclComm.get()), args.stream));

    CHK_RET(hcclStreamSynchronize(args.stream));

    // 3、DevtoHost
    CHK_RET(hrtMemSyncCopy(recvShape.data(), recvShape.size(), recvShapeMemPtr, recvShapeSize,
        HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_DEVICE_TO_HOST));

    HCCL_INFO("HcomGetRecvBeforeKernel get dynamic recv shape DONE and dim is [%lld]", recvShape[0]);

    return ge::GRAPH_SUCCESS;
}

ge::graphStatus LaunchRecvKernel(gert::KernelContext *context)
{
    uint64_t tensorSize = 0;
    std::vector<int64_t> recvShape(sizeof(int64_t) * SHAPE_INFO_COUNT, 0);

    SetLaunchKernelMode(true);

    // 获取算子输入输出资源
    HcomOpLaunchArgs launchArgs;
    if (GetRecvOpLaunchArgs(context, launchArgs) != HCCL_SUCCESS) {
        HCCL_ERROR("LaunchRecvKernel: get hcom op launch launchArgs failed.");
        return ge::GRAPH_FAILED;
    }

    // 获取通信域资源
    std::shared_ptr<hcclComm> hcclComm;
    if (HcomGetCommByGroup(launchArgs.opAttr.group, hcclComm)) {
        HCCL_ERROR("LaunchRecvKernel: get hcom comm handle failed. group:%s", launchArgs.opAttr.group);
        return ge::GRAPH_FAILED;
    }
    if (hcclComm->CreateCommCCLbuffer()) {
        HCCL_ERROR("LaunchRecvKernel: create ccl buffer failed. group:%s", launchArgs.opAttr.group);
        return ge::GRAPH_FAILED;
    }

    HCCL_DEBUG("Call LaunchRecvKernel, qosCfg value[%u]", launchArgs.opAttr.qosConfig);
    // 下发QosCofig给Runtime
    CHK_RET(hcclComm->SetQosCfg(launchArgs.opAttr.qosConfig));

    // 设置单算子mode
    HcclWorkflowMode lastWorkflowMode = GetWorkflowMode();
    if (SetWorkflowMode(HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) != HCCL_SUCCESS) {
        HCCL_ERROR("LaunchRecvKernel: config opbase mode failed");
        return ge::GRAPH_FAILED;
    }
    SetLaunchKernelMode(true);

    if (HcomGetRecvBeforeKernel(launchArgs, recvShape) != ge::GRAPH_SUCCESS) {
        HCCL_ERROR("LaunchRecvKernel: get recvShape fail");
        return ge::GRAPH_FAILED;
    }

    // 1. create && append outputshape
    auto outputShape = context->GetOutputPointer<gert::StorageShape>(0U);
    if (outputShape == nullptr) {
        return ge::GRAPH_FAILED;
    }
    const uint32_t dimNum = recvShape[0]; // recvshape n维的长度
    if (dimNum > gert::Shape::kMaxDimNum || dimNum <= 0) {
        HCCL_ERROR("dimNum=[%u] is greater than MaxDimNum[%zu] or less than or equal to 0", dimNum, gert::Shape::kMaxDimNum);
        return ge::GRAPH_FAILED;
    }
    gert::Shape& originShape = outputShape->MutableOriginShape();
    gert::Shape& storageShape = outputShape->MutableStorageShape();
    originShape.SetDimNum(0);
    storageShape.SetDimNum(0);
    for (size_t i = 1U; i <= dimNum; ++i) {
        originShape.AppendDim(static_cast<int64_t>(recvShape[i]));
        storageShape.AppendDim(static_cast<int64_t>(recvShape[i]));
    }

    // 2. create tensor_data
    ge::DataType geDataType = context->GetInputValue<ge::DataType>(static_cast<int32_t>(RecvOpLaunchInput::DATATYPE));
    ge::graphStatus getSizeRet = CalcAlignedSizeByShape(originShape, geDataType, tensorSize);
    if (getSizeRet != ge::GRAPH_SUCCESS) {
        HCCL_ERROR("LaunchRecvKernel: get tensor size failed");
        return ge::GRAPH_FAILED;
    }

    auto gertAllocator
            = context->GetInputValue<gert::GertAllocator *>(static_cast<int32_t>(RecvOpLaunchInput::ALLOCATOR));
    if (gertAllocator == nullptr) {
        HCCL_ERROR("LaunchRecvKernel: get gertAllocator failed");
        return ge::GRAPH_FAILED;
    }

    auto outputTensorData = context->GetOutputPointer<gert::GertTensorData>(1U);

    *outputTensorData = gertAllocator->MallocTensorData(tensorSize);

    SetLaunchKernelMode(false);

    // 调用算子kernel，实现算子下发
    uint64_t recvCount = (tensorSize - 32) / SIZE_TABLE[launchArgs.opAttr.dataType];
    if (HcclRecv(outputTensorData->GetAddr(), recvCount, launchArgs.opAttr.dataType, launchArgs.opAttr.op.recv.srcRank,
        static_cast<HcclComm>(hcclComm.get()), launchArgs.stream) != HCCL_SUCCESS) {
        SetLaunchKernelMode(false);
        return ge::GRAPH_FAILED;
    }

    SetLaunchKernelMode(false);
    HCCL_INFO("LaunchRecvKernel dispatch hcclrecv DONE");

    // workflowmode状态回置
    if (SetWorkflowMode(lastWorkflowMode) != HCCL_SUCCESS) {
        return ge::GRAPH_FAILED;
    }
    HCCL_INFO("LaunchRecvKernel dispatch DONE");
    return ge::GRAPH_SUCCESS;
}
REGISTER_KERNEL(LaunchRecvKernel).RunFunc(LaunchRecvKernel).OutputsCreator(BuildHcclOutputShapeOutputs);
}